import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import SGDRegressor
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error

# -------------------------------
# 1. 生成阿基米德螺线数据
# -------------------------------
a = 0.5  # 控制螺线间距
theta = np.linspace(0, 8 * np.pi, 500)  # 角度从0到8π
r = a * theta  # 理论上的极径 r = a * theta

# 添加噪声模拟真实观测数据
noise = np.random.normal(0, 0.5, size=r.shape)
r_noisy = r + noise

# 数据维度检查
print("原始 theta shape:", theta.shape)

# 转换为二维数组 (n_samples, n_features)
X = theta.reshape(-1, 1)  # 必须是二维数组
y = r_noisy

# -------------------------------
# 2. 标准化输入特征（防止梯度不稳定）
# -------------------------------
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)  # fit+transform训练集

# -------------------------------
# 3. 使用梯度下降法进行回归建模
# -------------------------------
model = SGDRegressor(max_iter=1000, tol=1e-5, random_state=42)
model.fit(X_scaled, y)

# 预测
y_pred = model.predict(X_scaled)

# -------------------------------
# 4. 计算MSE评估指标
# -------------------------------
mse = mean_squared_error(y, y_pred)
print(f"模型MSE: {mse:.4f}")

# -------------------------------
# 5. 极坐标可视化：突出显示螺旋线结构
# -------------------------------
fig = plt.figure(figsize=(12, 6))

# 子图1：真实螺旋线（含噪声）
ax1 = fig.add_subplot(1, 2, 1, projection='polar')
ax1.scatter(theta, r_noisy, color='blue', label='含噪声的真实数据', alpha=0.6)
ax1.plot(theta, r, color='red', linewidth=2, label='理论螺旋线')
ax1.set_title('真实阿基米德螺线')
ax1.legend(loc='upper right')

# 子图2：模型拟合结果
ax2 = fig.add_subplot(1, 2, 2, projection='polar')
ax2.scatter(theta, r_noisy, color='blue', label='含噪声的真实数据', alpha=0.6)
ax2.plot(theta, y_pred, color='green', linewidth=2, label=f'梯度下降拟合 (MSE={mse:.2f})')
ax2.set_title('模型拟合结果')
ax2.legend(loc='upper right')

plt.tight_layout()
plt.show()
